import cv2
import numpy as np

def order_points(pts):
    """对四边形的四个顶点进行排序，顺序为：左上、右上、右下、左下"""
    rect = np.zeros((4, 2), dtype="float32")
    
    # 计算顶点坐标的和，左上角的和最小，右下角的和最大
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    
    # 计算顶点坐标的差，右上角的差最小，左下角的差最大
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    
    return rect

def four_point_transform(image, pts):
    """对图像进行四点透视变换"""
    rect = order_points(pts)
    (tl, tr, br, bl) = rect
    
    # 计算变换后图像的宽度
    widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    maxWidth = max(int(widthA), int(widthB))
    
    # 计算变换后图像的高度
    heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
    heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
    maxHeight = max(int(heightA), int(heightB))
    
    # 构建变换后的图像的四个顶点坐标
    dst = np.array([
        [0, 0],
        [maxWidth - 1, 0],
        [maxWidth - 1, maxHeight - 1],
        [0, maxHeight - 1]], dtype="float32")
    
    # 计算透视变换矩阵并应用
    M = cv2.getPerspectiveTransform(rect, dst)
    warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
    
    return warped

def preprocess_image(image):
    """图像预处理：灰度化、降噪"""
    # 转换为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 高斯滤波降噪
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    return gray, blurred

def detect_edges(blurred):
    """边缘检测"""
    # 自适应阈值处理，增强文本边缘 - 调整参数以提高灵敏度
    thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                  cv2.THRESH_BINARY_INV, 11, 3)  # 减小C值提高灵敏度
    # Canny边缘检测 - 降低阈值以检测更多边缘
    edged = cv2.Canny(blurred, 30, 150, apertureSize=3)  # 降低阈值
    # 合并阈值和边缘结果，增强边缘
    combined_edges = cv2.bitwise_or(thresh, edged)
    
    # 形态学操作，闭运算连接断开的边缘 - 增大核大小
    kernel = np.ones((7, 7), np.uint8)  # 增大核大小
    closed_edges = cv2.morphologyEx(combined_edges, cv2.MORPH_CLOSE, kernel)
    
    return closed_edges

def find_document_contour(edges, image):
    """查找并筛选文档轮廓"""
    # 查找轮廓
    contours, _ = cv2.findContours(edges.copy(), cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    
    # 按轮廓面积降序排序
    contours = sorted(contours, key=cv2.contourArea, reverse=True)[:15]  # 增加考虑的轮廓数量
    
    document_contour = None
    
    # 遍历轮廓
    for contour in contours:
        # 计算轮廓周长
        perimeter = cv2.arcLength(contour, True)
        # 多边形逼近 - 调整精度参数
        approx = cv2.approxPolyDP(contour, 0.03 * perimeter, True)  # 增大精度系数
        
        # 如果轮廓有四个顶点或接近四个顶点，则认为找到了文档
        if 4 <= len(approx) <= 6:  # 放宽顶点数量限制
            # 如果不是4个点，尝试强制转为4个点
            if len(approx) != 4:
                rect = cv2.minAreaRect(approx)
                box = cv2.boxPoints(rect)
                approx = np.array(box, dtype=np.int32)  # 修改 np.int0 为 np.int32
            
            # 计算轮廓面积与图像面积的比例
            contour_area = cv2.contourArea(approx)
            image_area = image.shape[0] * image.shape[1]
            area_ratio = contour_area / image_area
            
            # 面积比例应在合理范围内 - 放宽范围
            if 0.05 < area_ratio < 0.95:  # 放宽面积比例限制
                document_contour = approx
                break
    
    # 如果没有找到四边形轮廓，尝试使用霍夫变换检测直线
    if document_contour is None:
        # 霍夫变换检测直线 - 降低阈值以检测更多线段
        lines = cv2.HoughLinesP(edges, 1, np.pi/180, threshold=80, minLineLength=80, maxLineGap=20)
        
        if lines is not None and len(lines) >= 4:
            # 根据霍夫变换结果构建四边形轮廓
            x_min, y_min = np.inf, np.inf
            x_max, y_max = 0, 0
            
            for line in lines:
                x1, y1, x2, y2 = line[0]
                x_min = min(x_min, x1, x2)
                y_min = min(y_min, y1, y2)
                x_max = max(x_max, x1, x2)
                y_max = max(y_max, y1, y2)
            
            # 确保坐标在图像范围内
            x_min = max(0, x_min)
            y_min = max(0, y_min)
            x_max = min(image.shape[1] - 1, x_max)
            y_max = min(image.shape[0] - 1, y_max)
            
            # 构建四边形轮廓
            document_contour = np.array([
                [[x_min, y_min]],
                [[x_max, y_min]],
                [[x_max, y_max]],
                [[x_min, y_max]]
            ], dtype=np.int32)
    
    # 添加额外的备选方案：如果仍然没有找到轮廓，尝试使用最大轮廓
    if document_contour is None and contours:
        largest_contour = contours[0]  # 最大的轮廓
        # 使用最小面积矩形拟合
        rect = cv2.minAreaRect(largest_contour)
        box = cv2.boxPoints(rect)
        document_contour = np.array(box, dtype=np.int32)  # 修改 np.int0 为 np.int32
    
    return document_contour

def enhance_document(warped):
    """增强扫描后的文档图像"""
    # 转换为灰度图
    gray = cv2.cvtColor(warped, cv2.COLOR_BGR2GRAY)
    
    # 应用双边滤波，保留边缘的同时去除噪声
    filtered = cv2.bilateralFilter(gray, 9, 75, 75)
    
    # 应用CLAHE（对比度受限的自适应直方图均衡化）增强对比度
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced_contrast = clahe.apply(filtered)
    
    # 使用更温和的自适应阈值处理，增强文本
    # 增大blockSize和减小C值，使文本更清晰
    thresh = cv2.adaptiveThreshold(enhanced_contrast, 255, 
                                  cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                  cv2.THRESH_BINARY, 15, 5)
    
    # 创建彩色输出图像
    result = cv2.cvtColor(thresh, cv2.COLOR_GRAY2BGR)
    
    # 保留原始彩色信息的版本（可选）
    # 创建一个掩码，其中文本区域为白色
    _, text_mask = cv2.threshold(thresh, 127, 255, cv2.THRESH_BINARY_INV)
    
    # 将原始彩色图像与增强后的图像混合
    # 这样可以保留彩色信息，同时增强文本清晰度
    color_enhanced = warped.copy()
    color_enhanced[text_mask == 255] = [0, 0, 0]  # 文本区域设为黑色
    
    # 根据图像亮度决定返回黑白版本还是彩色版本
    # 计算原始图像的平均亮度
    avg_brightness = np.mean(gray)
    
    # 如果图像较暗或对比度较低，返回黑白版本
    if avg_brightness < 120:
        return result
    else:
        # 否则返回彩色增强版本
        return color_enhanced

def main():
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    
    # 检查摄像头是否成功打开
    if not cap.isOpened():
        print("无法打开摄像头")
        return
    
    # 获取摄像头分辨率
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # 创建窗口
    cv2.namedWindow("文档扫描", cv2.WINDOW_NORMAL)
    
    # 添加调试模式标志
    debug_mode = False
    
    while True:
        # 读取一帧
        ret, frame = cap.read()
        if not ret:
            print("无法获取视频帧")
            break
        
        # 复制原始帧用于显示
        original = frame.copy()
        
        # 图像预处理
        gray, blurred = preprocess_image(frame)
        
        # 边缘检测
        edges = detect_edges(blurred)
        
        # 查找文档轮廓
        document_contour = find_document_contour(edges, frame)
        
        # 创建结果显示图像
        result = np.zeros((frame_height, frame_width * 2, 3), dtype=np.uint8)
        
        # 左侧显示原始图像
        result[:, :frame_width] = original
        
        # 如果找到文档轮廓
        if document_contour is not None:
            # 在原始图像上绘制文档轮廓
            cv2.drawContours(original, [document_contour], -1, (0, 255, 0), 2)
            result[:, :frame_width] = original
            
            # 透视变换
            warped = four_point_transform(frame, document_contour.reshape(4, 2))
            
            # 增强文档图像
            enhanced = enhance_document(warped)
            
            # 调整大小以适应显示区域
            enhanced_resized = cv2.resize(enhanced, (frame_width, frame_height))
            
            # 右侧显示处理后的文档
            result[:, frame_width:] = enhanced_resized
            
            # 显示文档检测状态
            cv2.putText(result, "文档已检测", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                        1, (0, 255, 0), 2)
        else:
            # 如果没有找到文档，右侧显示边缘检测结果或原始图像
            if debug_mode:
                edges_colored = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
                result[:, frame_width:] = edges_colored
            else:
                result[:, frame_width:] = frame
            
            # 显示文档检测状态
            cv2.putText(result, "未检测到文档", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                        1, (0, 0, 255), 2)
        
        # 显示结果
        cv2.imshow("文档扫描", result)
        
        # 按键处理
        key = cv2.waitKey(1) & 0xFF
        
        # 按'q'键退出
        if key == ord('q'):
            break
        
        # 按'd'键切换调试模式
        elif key == ord('d'):
            debug_mode = not debug_mode
            print(f"调试模式: {'开启' if debug_mode else '关闭'}")
        
        # 按's'键保存扫描结果
        elif key == ord('s') and document_contour is not None:
            # 保存增强后的文档图像
            cv2.imwrite("d:\\code\\video-scan\\scanned_document.jpg", enhanced)
            # 保存原始扫描区域图像（未经增强处理）
            cv2.imwrite("d:\\code\\video-scan\\original_document.jpg", warped)
            print("扫描结果已保存 - 增强版本和原始版本")
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()